#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Author: '习惯'
# @Email: 'songbing513@gmail.com'
# @Date: '2025/9/1 21:57'
# @Project: 'McpAgent'
# @File: '02.py'
# @Software: 'PyCharm'

from typing import Optional
import json
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field, ValidationError


"""
使用LangChain的ChatOpenAI类，并配置为使用DeepSeek模型。
我们定义了一个Pydantic模型Joke，
它有三个字段：setup（笑话的开头）、punchline（笑话的包袱）和rating（笑话的有趣程度评分，可选，范围0-10）。

DeepSeek不支持:
    with_structured_output 方法.
    response_format 结构化输出功能
    但是可以通过自定义输出解析器实现结构化输出
    详细参考: https://github.com/huggingface/transformers/blob/main/src/transformers/pipelines/text2text_generation.py#L114
修改提示词，要求模型返回一个符合特定JSON格式的字符串
手动解析返回的JSON字符串为Joke对象。
注意：由于模型可能返回非标准JSON，我们需要添加错误处理。

chatGPT 写法
class Joke(BaseModel):
    setup: str = Field(description="The setup of the joke.")
    punchline: str = Field(description="The punchline of the joke.")
    rating: Optional[int] = Field(description="How funny the joke is, on a scale from 0 to 10.")
response = llm.with_structured_output("Tell me a joke about cats")
response.invoke("Tell me joke about cats")
"""

# 初始化 DeepSeek 模型
llm = ChatOpenAI(
    model="deepseek-chat",
    temperature=0,
    base_url="https://api.deepseek.com/v1",
    api_key="sk-3291e6df463f408eaeea2629bd84fc6a",
)


# 定义笑话结构
class Joke(BaseModel):
    """Joke to tell user."""
    setup: str = Field(description="The setup of the joke.")
    punchline: str = Field(description="The punchline of the joke.")
    rating: Optional[int] = Field(description="How funny the joke is, on a scale from 0 to 10.")


def get_joke(topic: str) -> Joke:
    """获取关于指定主题的笑话"""
    # 构建明确的JSON格式提示
    prompt = f"""
    请创建一个关于{topic}的笑话，并以严格的JSON格式返回，包含以下字段：
    - "setup": 笑话的开场白
    - "punchline": 笑话的包袱/笑点
    - "rating": 笑话的有趣程度评分（0-10分）

    示例格式：
    {{
        "setup": "为什么程序员讨厌自然？",
        "punchline": "因为那里有太多的虫子！",
        "rating": 8
    }}

    现在请创建关于{topic}的笑话：
    """

    # 调用模型获取响应
    response = llm.invoke(prompt)

    try:
        # 尝试从响应中提取JSON内容
        content = response.content.strip()

        # 处理可能的代码块标记
        if content.startswith("```json"):
            content = content.split("```json")[1].split("```")[0].strip()
        elif content.startswith("```"):
            content = content.split("```")[1].split("```")[0].strip()

        # 解析JSON并验证
        joke_data = json.loads(content)
        return Joke(**joke_data)

    except (json.JSONDecodeError, ValidationError) as e:
        # 如果解析失败，使用更简单的回退方法
        print(f"JSON解析失败，使用回退方法: {e}")
        return parse_joke_from_text(response.content)


def parse_joke_from_text(text: str) -> Joke:
    """从文本中解析笑话结构"""
    # 尝试从文本中提取结构
    setup = ""
    punchline = ""
    rating = None

    # 简单解析逻辑
    lines = text.split('\n')
    for line in lines:
        if "setup" in line.lower() or "开场" in line.lower():
            setup = line.split(':', 1)[-1].strip()
        elif "punchline" in line.lower() or "包袱" in line.lower() or "笑点" in line.lower():
            punchline = line.split(':', 1)[-1].strip()
        elif "rating" in line.lower() or "评分" in line.lower():
            try:
                rating_part = line.split(':', 1)[-1].strip()
                rating = int(''.join(filter(str.isdigit, rating_part)))
            except (ValueError, IndexError):
                pass

    # 确保必要字段存在
    if not setup or not punchline:
        # 如果无法解析，创建默认笑话
        setup = "为什么这个笑话无法解析？"
        punchline = "因为API不支持结构化输出！"
        rating = 3

    return Joke(setup=setup, punchline=punchline, rating=rating)


# 获取并打印笑话
try:
    joke = get_joke("cats")
    print("\n生成的猫咪笑话:")
    print(f"开场: {joke.setup}")
    print(f"笑点: {joke.punchline}")
    print(f"评分: {joke.rating}/10")
except Exception as e:
    print(f"获取笑话时出错: {e}")
    # 创建默认笑话作为回退
    fallback_joke = Joke(
        setup="为什么API不支持结构化输出？",
        punchline="因为它还在学习如何讲笑话！",
        rating=5
    )
    print("\n备用笑话:")
    print(f"开场: {fallback_joke.setup}")
    print(f"笑点: {fallback_joke.punchline}")
    print(f"评分: {fallback_joke.rating}/10")